import math
from contextlib import asynccontextmanager
from typing import (Type, Generic, Any, Sequence, Self, Optional, Dict, Tuple, AsyncIterator, Callable, Awaitable,
                    Iterable)

from sqlalchemy import select, SelectBase, and_, func, exc, or_
from sqlalchemy.ext.asyncio.session import AsyncSession
from sqlalchemy.orm import selectinload, with_loader_criteria, joinedload
from sqlalchemy.sql.selectable import ForUpdateArg

from app.core.custom_exception import ConflictException
from app.core.generic_types import SQLModelType
from app.models.base import IdType
from app.schemas.query_schema import BaseQueryIn, OperatorsEnum, ConditionEnum
from app.schemas.result_schema import PageModel
from app.utils.list_extension import distinct


class AsyncSessionExtension(Generic[SQLModelType]):
    def __init__(self, session: AsyncSession, model: Optional[Type[SQLModelType]]=None):
        self._model = model
        self._session = session
        self._filters = []
        self._includes = []
        self._order_by = []
        self._offset = None
        self._limit = None
        self._distinct = False
        self._current_loader = None  # 跟踪当前加载器，用于预先加载加载器的链式调用
        self._load_path =[] # 记录加载顺序
        self._select_many_load_path =[] # 记录显示加载顺序

    @property
    def session(self) -> AsyncSession:
        """验证 AsyncSession 是否已设置"""
        if self._session is None:
            raise ValueError("AsyncSession未注入！请通过依赖注入设置session")
        return self._session

    @property
    def model(self) -> Type[SQLModelType]:
        """验证模型是否已设置"""
        if self._model is None:
            raise ValueError("模型未设置！请先调用set(Entity)方法。")
        return self._model

    #region 构建查询语句


    def set(self, model: Type[SQLModelType]) -> Self:
        """
        设置模型
        :param model: 要查询的模型
        :return: Self
        """
        self._model = model
        return self

    def where(self, *filters: Any) -> Self:
        """
        添加查询条件
        :param filters: 查询条件
        :return: Self
        """
        self._filters.extend(filters)
        return self

    def filter_by(self,**kwargs: Any) -> Self:
        """
        根据键值对添加查询条件
        :param kwargs: 键值对参数
        :return: Self
        """
        for field_name, value in kwargs.items():
            if hasattr(self.model, field_name):
                self._filters.append(getattr(self.model, field_name) == value)
            else:
                raise ValueError(f"名称为{self.model.__name__}的模型中，不存在字段名为{field_name}的属性")
        return self

    def dynamic_query(self, query_in: BaseQueryIn) -> Self:
        """
        根据前端传入的动态查询模型构建查询条件
        :param query_in: 动态查询模型
        :return: Self
        """
        if query_in.query_conditions:
            filters = []

            for qc in query_in.query_conditions:
                if not hasattr(self.model, qc.field):
                    continue

                field = getattr(self.model, qc.field)
                value = qc.value
                second_value = qc.second_value

                # 如果参数名称包含或者等于id,则将其转换成uuid类型
                if qc.field.endswith("_id") or qc.field == "id":
                    value = IdType(value)

                match qc.operator:
                    # 没有，不执行任何操作
                    case OperatorsEnum.none:
                        break
                    # 等于 =
                    case OperatorsEnum.equal:
                        filters.append(field == value)
                        break
                    # 不等于 !=
                    case OperatorsEnum.not_equal:
                        filters.append(field != value)
                        break
                    # 大于 >
                    case OperatorsEnum.greater_than:
                        filters.append(field > value)
                        break
                    # 大于等于 >=
                    case OperatorsEnum.greater_than_or_equal_to:
                        filters.append(field >= value)
                        break
                    # 小于 <
                    case OperatorsEnum.less_than:
                        filters.append(field < value)
                        break
                    # 小于等于 <=
                    case OperatorsEnum.less_than_or_equal_to:
                        filters.append(field <= value)
                        break
                    # 包含
                    case OperatorsEnum.contains:
                        filters.append(field.like(f"%{value}%"))
                        break
                    # 不包含
                    case OperatorsEnum.not_contains:
                        filters.append(~field.like(f"%{value}%"))
                        break
                    # 以什么开始
                    case OperatorsEnum.starts_with:
                        filters.append(field.like(f"{value}%"))
                        break
                    # 以什么结束
                    case OperatorsEnum.ends_with:
                        filters.append(field.like(f"%{value}"))
                        break
                    # 在范围以内
                    case OperatorsEnum.range:
                        if second_value is not None:
                            filters.append(and_(field >= value, field <= second_value))
                        break
                    # 在范围以外
                    case OperatorsEnum.out_range:
                        if second_value is not None:
                            filters.append(or_(field < value, field > second_value))
                        break

            if filters:
                # 判断条件 而且
                if query_in.condition_relation == ConditionEnum.and_also:
                    self._filters.append(and_(*filters))
                else:
                    # 或者
                    self._filters.append(or_(*filters))

            if query_in.sort and query_in.order:
                self.dynamic_order_by(query_in.sort, query_in.order)

        return self

    def dynamic_order_by(self, field_name: str, order:bool) -> Self:
        """
        动态查询排序
        :param field_name: 排序字段名
        :param order: 正序还是倒序
        :return: Self
        """
        if hasattr(self.model, field_name):
            sort_field = getattr(self.model, field_name)
            self._order_by.append(sort_field.desc() if not order else sort_field.asc())
        return self
    def _chainload(self, relationship: Any, *, recursion_depth: Optional[int] = None, criteria: Optional[Tuple[Type[Any], Any]] = None) -> Self:
        """
        使用 SELECT IN 预先加载来加载给定的属性
        :param relationship: 关系
        :param recursion_depth: 加载深度，仅用于自身循环引用
        :param criteria: 当前映射的过滤条件
        :return: Self
        """
        if self._current_loader is None:
            # 创建新加载器
            self._current_loader = selectinload(relationship, recursion_depth=recursion_depth)
            self._includes.append(self._current_loader)
        else:
            # 添加到当前加载器
            self._current_loader = self._current_loader.selectinload(relationship, recursion_depth=recursion_depth)
            self._includes[0] = self._current_loader
        if criteria:
            model_cls, condition = criteria
            self._includes.append(with_loader_criteria(model_cls, condition))
        return self

    def _not_chainload(self, relationship: Any, *, recursion_depth: Optional[int] = None, criteria: Optional[Tuple[Type[Any], Any]] = None) -> Self:
        """
        使用 SELECT IN 预先加载来加载给定的属性
        :param relationship: 关系
        :param recursion_depth: 加载深度，仅用于自身循环引用
        :param criteria: 当前映射的过滤条件
        :return: Self
        """
        # 创建新加载器
        self._current_loader = selectinload(relationship, recursion_depth=recursion_depth)
        self._includes.append(self._current_loader)
        if criteria:
            model_cls, condition = criteria
            self._includes.append(with_loader_criteria(model_cls, condition))
        return self

    def _joinedload(self,relationship: Any,*,criteria: Optional[Tuple[Type[Any], Any]] = None, **kw: Any,) -> Self:
        """
        使用 join 连接的预先加载来加载给定的属性
        :param relationship: 关系
        :param criteria: 当前映射的过滤条件
        :param kw: innerjoin 如果 True，则指示连接的急切加载应使用内连接，而不是默认的左外连接
        :return: Self
        """
        if self._current_loader is None:
            # 创建新加载器
            self._current_loader = joinedload(relationship, **kw)
            self._includes.append(self._current_loader)
        else:
            # 添加到当前加载器
            self._current_loader = self._current_loader.joinedload(relationship, **kw)
            self._includes[0] = self._current_loader
        if criteria:
            model_cls, condition = criteria
            self._includes.append(with_loader_criteria(model_cls, condition))
        return self

    def join(self,relationship: Any,*,criteria: Optional[Tuple[Type[Any], Any]] = None, **kw: Any,) -> Self:
        """
        使用 join 连接的预先加载来加载给定的属性
        :param relationship: 关系
        :param criteria: 当前映射的过滤条件
        :param kw: innerjoin 如果 True，则指示连接的急切加载应使用内连接，而不是默认的左外连接
        :return: Self
        """
        return self._joinedload(relationship,criteria=criteria, **kw)

    def include(self, relationship: Any,*, recursion_depth: Optional[int] = None, criteria: Optional[Tuple[Type[Any], Any]] = None) -> Self:
        """
        使用 SELECT IN 预先加载来加载给定的属性
        :param relationship: 关系
        :param recursion_depth: 加载深度，仅用于自身循环引用
        :param criteria: 当前映射的过滤条件
        :return: Self
        """
        if self._current_loader is None:
            self._load_path.append(relationship)
            return self._chainload(relationship, recursion_depth=recursion_depth, criteria=criteria)
        else:
            last_path = self._load_path[-1]
            self._load_path.append(relationship)
            if hasattr(last_path,relationship.key) :
                return self._chainload(relationship, recursion_depth=recursion_depth, criteria=criteria)
            else:
                return self._not_chainload(relationship, recursion_depth=recursion_depth, criteria=criteria)



    def select_many(self, relationship: Any,*, recursion_depth: Optional[int] = None, criteria: Optional[Tuple[Type[Any], Any]] = None) -> Self:
        """
        使用 SELECT IN 显示加载来加载给定的属性
        :param relationship: 关系
        :param recursion_depth: 加载深度，仅用于自身循环引用
        :param criteria: 当前映射的过滤条件
        :return: Self
        """
        self._select_many_load_path.append(relationship)
        return self._chainload(relationship, recursion_depth=recursion_depth, criteria=criteria)

    def order_by(self,*order_by:Any)->Self:
        """
        添加排序条件
        :param order_by: 排序条件
        :return: Self
        """
        self._order_by.extend(order_by)
        return self

    def limit(self, limit: int) ->Self:
        """
        限制返回记录数
        :param limit: 最大返回记录数
        :return: Self
        """
        self._limit = limit
        return self

    def offset(self, offset: int)->Self:
        """
        设置偏移量
        :param offset: 偏移量
        :return: Self
        """
        self._offset = offset
        return self

    def distinct(self) ->Self:
        """
        去重查询
        :return: Self
        """
        self._distinct = True
        return self

    @property
    def not_deleted_condition(self)->bool:
        """
        获取未删除数据的查询条件
        :return: bool
        """
        return self.model.is_delete == False

    @property
    def not_deleted(self) -> SelectBase[SQLModelType]:
        """
        获取未删除数据的查询
        :return: SelectBase[ModelType]
        """
        return select(self.model).where(self.not_deleted_condition)

    async def _reset_query(self) -> None:
        """重置查询条件"""
        self._filters = []
        self._includes = []
        self._order_by = []
        self._limit = None
        self._offset = None
        self._distinct = False
        self._current_loader = None  # 跟踪当前加载器
        self._load_path = []

    def _build_query(self) -> SelectBase[SQLModelType]:
        """构建查询语句"""
        stmt = self.not_deleted

        if self._filters:
            stmt = stmt.where(and_(*self._filters))

        # 处理预先加载选项
        if self._includes:
            stmt = stmt.options(*self._includes)

        if self._order_by:
            stmt = stmt.order_by(*self._order_by)

        if self._limit is not None:
            stmt = stmt.limit(self._limit)

        if self._offset is not None:
            stmt = stmt.offset(self._offset)

        if self._distinct:
            stmt = stmt.distinct()

        return stmt

    #endregion

    # region 查询
    async def _count(self, page_size: int = 100) -> Tuple[int, int]:
        """
        获取总数
        :param page_size: 分页大小
        :return: 数据总数量和分页页数
        """
        stmt = select(func.count()).select_from(self.not_deleted)
        if self._filters:
            stmt = stmt.where(and_(*self._filters))
        total = (await self.session.execute(stmt)).scalar()
        total_pages = math.ceil(total / page_size)
        return total, total_pages

    async def count(self) -> int:
        """
        统计记录数量
        :return: 数据总数量
        """
        stmt = select(func.count()).select_from(self.not_deleted)
        if self._filters:
            stmt = stmt.where(and_(*self._filters))
        result = await self.session.execute(stmt)
        count = result.scalar_one()
        await self._reset_query()
        return count

    async def _to_list(self) -> Sequence[SQLModelType]:
        """
        执行查询语句，获得所有符合条件的数据
        :return: 所有符合条件的数据序列
        """
        stmt = self._build_query()
        result = (await self.session.execute(stmt)).scalars().all()
        await  self._reset_query()
        return result

    async def _to_last_load_list(self) -> Sequence[SQLModelType]:
        """
        用于将显示加载的最后一个级的数据组合成一个序列
        :return: 最后一个级的数据组合成的序列
        :raise: 没有调用 select_many 抛出异常
        """
        result = await self._to_list()
        last_load_list = []
        if self._select_many_load_path and len(self._select_many_load_path) > 1:
            second_to_last_load = self._select_many_load_path[-2]
            last_load = self._select_many_load_path[-1]
            for item1 in result:
                item1_relationship = getattr(item1,second_to_last_load.key)
                if item1_relationship is not None and isinstance(item1_relationship, list):
                    for item2 in item1_relationship:
                        item2_relationship = getattr(item2,last_load.key)
                        if item2_relationship is not None and isinstance(item2_relationship, list):
                            last_load_list.extend(item2_relationship)
        elif self._select_many_load_path and len(self._select_many_load_path) == 1:
            last_load = self._select_many_load_path[-1]
            for item in result:
                relationship = getattr(item,last_load.key)
                if relationship is not None and isinstance(relationship, list):
                    last_load_list.extend(relationship)
        self._select_many_load_path = []
        return distinct(last_load_list)

    async def to_list(self) -> Sequence[SQLModelType]:
        """
        执行查询语句，获得所有符合条件的数据
        :return: 所有符合条件的数据序列
        """
        if len(self._select_many_load_path) == 0:
            return await self._to_list()
        else:
            return await self._to_last_load_list()

    async def to_dict_list(self) -> Sequence[Dict[str, Any]]:
        """
        执行查询语句，获得所有符合条件的数据，并将其转为字典序列
        :return: 所有符合条件的字典数据序列
        """
        stmt = self._build_query()
        result = await self.session.execute(stmt)
        dict_list = [row.dict() for row in result.mappings().all()]
        await  self._reset_query()
        return dict_list

    async def to_page_list(self, current_page: int = 1, page_size: int = 100) -> PageModel[SQLModelType]:
        """
        执行查询语句，分页查询获得所有符合条件的数据
        :param current_page: 当前页码
        :param page_size: 每一页的大小
        :return:
        """
        self._offset = (current_page - 1) * page_size
        self._limit = page_size

        total,total_pages =await self._count(page_size=page_size)

        stmt = self._build_query()
        result = await self.session.execute(stmt)
        items = result.scalars().all()

        await  self._reset_query()
        return PageModel(items=items,total=total,total_pages=total_pages,current_page=current_page,page_size=page_size)


    async def first(self) -> Optional[SQLModelType]:
        """
        执行查询语句，获取第一条记录
        :return: 获取的第一条记录
        """
        stmt = self._build_query().limit(1)
        result = (await self.session.execute(stmt)).scalars().first()
        await  self._reset_query()
        return result

    async def one(self) -> SQLModelType:
        """
        执行查询语句，获取单条记录
        :return: 获取的单条记录
        :raise: ValueError 如果不存在或多条则抛出异常
        """
        stmt = self._build_query()
        result = await self.session.execute(stmt)
        try:
            return result.scalars().one()
        except Exception as e:
            await  self._reset_query()
            raise ValueError(f"预期只有一个结果，查询时没有或有多个结果: {stmt}") from e
        finally:
            await  self._reset_query()

    async def one_or_none(self) -> Optional[SQLModelType]:
        """
         执行查询语句，获取单条记录
         :return: 获取的单条记录,不存在则返回None
         :raise: ValueError 多条则抛出异常
         """
        stmt = self._build_query()
        try:
            result = (await self.session.execute(stmt)).scalars().one_or_none()
            return result
        except exc.MultipleResultsFound as e:
            await  self._reset_query()
            raise ValueError(f"预期只有一个结果，查询时没有或有多个结果: {stmt}") from e
        finally:
            await  self._reset_query()

    async def scalar(self) -> Any:
        """获取标量结果"""
        stmt = self._build_query()
        result = (await self.session.execute(stmt)).scalar()
        await  self._reset_query()
        return result

    async def scalars(self) -> Sequence[SQLModelType]:
        """获取标量结果序列"""
        stmt = self._build_query()
        result = (await self.session.execute(stmt)).scalars().all()
        await self._reset_query()
        return result

    async def scalar_one_or_none(self)  -> Optional[SQLModelType]:
        stmt = self._build_query()
        result = (await self.session.execute(stmt)).scalar_one_or_none()
        await  self._reset_query()
        return result
    #endregion

    #region 增、删、改

    async def save(self, entity: SQLModelType, *, flush: bool = False, refresh:bool=False) -> Optional[SQLModelType]:
        """
        提交单个记录
        :param entity: 要提交的实体
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 提交成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        try:
            self.session.add(entity)
            if flush:
                await self.session.flush()
            await self.session.commit()
            if refresh:
                await self.session.refresh(entity)
            return entity
        except exc.IntegrityError as e:
            await self.session.rollback()
            raise ConflictException(message=f"发生冲突：{self.model.__name__}: {str(e)}")


    async def save_all(self, entities: Sequence[SQLModelType], *, flush: bool = False, refresh:bool=False) -> Sequence[SQLModelType]:
        """
        批量提交记录
        :param entities: 要提交的实体列表
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 提交成功，返回实体模型列表
        :raise HTTPException 发生冲突，抛出异常
        """
        try:
            self.session.add_all(entities)
            if flush:
                await self.session.flush()
            else:
                await self.session.commit()
            if refresh:
                for entity in entities:
                    await self.session.refresh(entity)
            return entities
        except exc.IntegrityError as e:
            await self.session.rollback()
            raise ConflictException(message=f"发生冲突：{self.model.__name__}: {str(e)}")

    async def add(self, entity: SQLModelType, *, flush: bool = False, refresh:bool=False) -> Optional[SQLModelType]:
        """
        添加单个记录
        :param entity: 要添加的实体
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 添加成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        return await self.save(entity, flush=flush,refresh=refresh)

    async def add_all(self, entities: Sequence[SQLModelType], *, flush: bool = False, refresh:bool=False) -> Sequence[SQLModelType]:
        """
        批量添加记录
        :param entities: 要添加的实体列表
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 添加成功，返回实体模型列表
        :raise HTTPException 发生冲突，抛出异常
        """
        return await self.save_all(entities, flush=flush,refresh=refresh)

    async def bulk_insert(self, entities: Sequence[SQLModelType], *, return_defaults: bool = False, render_nulls: bool = False) -> None:
        """高性能批量插入"""
        await self.session.run_sync(
            lambda sync_session: sync_session.bulk_save_objects(
                entities,
                return_defaults=return_defaults,
                render_nulls=render_nulls
            )
        )

    async def bulk_update(self, entities: Sequence[SQLModelType]) -> None:
        """高性能批量更新"""
        await self.session.run_sync(
            lambda sync_session: sync_session.bulk_update_mappings(
                self.model,
                [entity.model_dump(exclude_unset=True) for entity in entities]
            )
        )

    async def update(self, entity: SQLModelType, *, flush: bool = False, refresh:bool=False) -> Optional[SQLModelType]:
        """
        更新单个记录
        :param entity: 要更新的实体
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 更新成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        return await self.save(entity, flush=flush,refresh=refresh)

    async def update_field(self, entity: SQLModelType, field_name:str, field_value:Any, *, flush: bool = False, refresh:bool=False) -> Optional[SQLModelType]:
        """
        更新单个记录的某个字段
        :param entity: 要更新的实体
        :param field_value: 字段名
        :param field_name:字段值
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 更新成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        if field_name in entity:
            entity[field_name] = field_value
        return await self.save( entity, flush=flush,refresh=refresh)

    async def update_all(self, entities: Sequence[SQLModelType], *, flush: bool = False, refresh:bool=False) ->  Sequence[SQLModelType]:
        """
        批量添加记录
        :param entities: 要添加的实体列表
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 添加成功，返回实体模型列表
        :raise HTTPException 发生冲突，抛出异常
        """
        return await self.save_all(entities, flush=flush,refresh=refresh)

    async def delete(self, entity: SQLModelType, *, flush: bool = False, refresh:bool=False) -> Optional[SQLModelType]:
        """
        删除单个记录
        :param entity: 要删除的实体
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 删除成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        entity.is_delete = True
        return await self.save( entity, flush=flush,refresh=refresh)

    async def delete_all(self, entities: Sequence[SQLModelType], *, flush: bool = False, refresh:bool=False) -> Sequence[SQLModelType]:
        """
        批量删除记录
        :param entities: 要删除的实体列表
        :param flush: 是否刷新数据库会话
        :param refresh: 是否刷新实体
        :return: 批量删除成功，返回实体模型
        :raise HTTPException 发生冲突，抛出异常
        """
        for entity in entities:
            entity.is_delete = True
        return await self.save_all( entities, flush=flush,refresh=refresh)

    #endregion

    # region 其他
    @asynccontextmanager
    async def nested_transaction(self):
        """开启嵌套事务"""
        async with self.session.begin_nested() as transaction:
            try:
                yield transaction
                await transaction.commit() # 提交嵌套事务
            except Exception:
                await transaction.rollback()  # 异常时回滚
                raise Exception("操作执行失败，嵌套事务已自动回滚")

    @asynccontextmanager
    async def transaction(self) -> AsyncIterator[None]:
        """支持async with语法的事务管理"""
        if self.session.in_transaction():
            yield
        else:
            async with self.session.begin():
                yield

    async def execute_in_transaction(self,operations: Sequence[Callable[[AsyncSession], Awaitable[Any]]]) -> Sequence[Any]:
        """在事务中执行多个操作"""
        results = []
        async with self.transaction():
            for op in operations:
                results.append(await op(self.session))
        return results

    async def execute(self, stmt: SelectBase[SQLModelType], *, reset_query: bool = True) -> Any:
        """执行原生SQL语句"""
        try:
            result = await self.session.execute(stmt)
            return result
        finally:
            if reset_query:
                await self._reset_query()

    async def flush(self) -> None:
        """刷新会话"""
        await self.session.flush()

    async def refresh(self, entity: SQLModelType, attribute_names: Iterable[str] | None = None, with_for_update: ForUpdateArg | None | bool | dict[str, Any] = None) -> None:
        """刷新实体"""
        await self.session.refresh(entity,attribute_names,with_for_update)

    async def refresh_all(self, entities: Sequence[SQLModelType], attribute_names: Iterable[str] | None = None,
                          with_for_update: ForUpdateArg | None | bool | dict[str, Any] = None) -> None:
        """批量刷新所有"""
        for entity in entities:
            await self.session.refresh(entity,attribute_names,with_for_update)

    async def close(self) -> None:
        # 确保会话被关闭
        if hasattr(self, '_session') and self._session is not None:
            await self._session.close()
    # endregion